from .a1_file_version_control_system import result_vcs
from .a3_netstructure import VisualRoute, AuditoryRoute, Brain
from QfUtil.TrainTool.train_classifier import train
from . import visual_model_name, auditory_model_name, brain_model_name, node_sizes


def train_inferior(train_loader, test_loader, turns_rate=1, visual_net=None, auditory_net=None):
    # model init
    if visual_net is None:
        visual_net = VisualRoute().cuda()
    if auditory_net is None:
        auditory_net = AuditoryRoute().cuda()

    # train visual
    def visaul_forward(net, din):
        vdin, adin = din
        return net(vdin)

    visual_train_result = train(train_loader=train_loader, test_loader=test_loader, network=visual_net,
                                turns=10 * turns_rate,
                                forward=visaul_forward)
    print('visual route train finished')

    # train auditory
    def auditory_forward(net, din):
        vdin, adin = din
        return net(adin)

    auditory_train_result = train(train_loader=train_loader, test_loader=test_loader, turns=10 * turns_rate,
                                  network=auditory_net,
                                  forward=auditory_forward)
    print('auditory route train finished')

    return visual_train_result, auditory_train_result


def get_trained_inferior(train_loader=None, test_loader=None, version=None, turns_rate=1):
    """
    loaders must not None if not specified version.

    :param train_loader:
    :param test_loader:
    :param version:
    :param turns_rate:
    :return:
    """

    # save path generate
    visual_route_model_path = result_vcs.generate_path(visual_model_name, version)
    auditory_route_model_path = result_vcs.generate_path(auditory_model_name, version)

    # if version specified, load and return
    if version is not None:
        visual_net = VisualRoute().load(visual_route_model_path)
        auditory_net = AuditoryRoute().load(auditory_route_model_path)
        return visual_net, auditory_net

    # models init
    visual_net = VisualRoute().cuda()
    auditory_net = AuditoryRoute().cuda()

    # train
    inferior_training_result = train_inferior(train_loader, test_loader, turns_rate, visual_net, auditory_net)

    # save
    visual_net.save(visual_route_model_path)
    auditory_net.save(auditory_route_model_path)
    result_vcs.save(inferior_training_result, 'inferior_training.result')

    return visual_net, auditory_net


def train_superior(train_loader, test_loader, turns_rate=1, brain=None):
    def forward(net, din):
        return brain(din)  # brain.forward(din)

    brain_train_result = train(train_loader=train_loader, test_loader=test_loader, turns=10 * turns_rate,
                               network=brain.cortex,
                               forward=forward)

    return brain_train_result


def get_trained_brain(train_loader=None, test_loader=None, inferior_version=None, version=None, turns_rate=10, depth=1,
                      node_size=200):
    # save path generate
    brain_model_path = result_vcs.generate_path(
        brain_model_name[:-6] + '_d' + str(depth) + '_s' + str(node_size) + '.model', version)

    # if version specified, load and return
    if version is not None:
        brain = Brain(depth=depth, node_size=node_size).load(brain_model_path).cuda()
        return brain

    # models init
    visual_net, auditory_net = get_trained_inferior(version=inferior_version)
    brain = Brain(visual_net, auditory_net, depth, node_size).cuda()

    # train
    superior_training_result = train_superior(train_loader, test_loader, turns_rate, brain)

    # save
    brain.save(brain_model_path)
    result_vcs.save(superior_training_result, 'superior_training.result')

    return brain


def get_series_trained_brain(train_loader=None, test_loader=None, inferior_version=None, version=None, turns_rate=10,
                             depth=1,
                             node_sizes=node_sizes):
    brains = {}
    for node_size in node_sizes:
        brains[node_size] = get_trained_brain(train_loader, test_loader, inferior_version, version, turns_rate, depth,
                                              node_size)
    return brains


def yield_series_trained_brain(train_loader=None, test_loader=None, inferior_version=None, version=None, turns_rate=10,
                               depth=1,
                               node_sizes=node_sizes):
    for node_size in node_sizes:
        brain = get_trained_brain(train_loader, test_loader, inferior_version, version, turns_rate, depth,
                                  node_size)
        yield node_size, brain
    return
